{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 机器学习预测模型\n",
    "\n",
    "本笔记本实现基于机器学习的股价预测模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import akshare as ak\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from sklearn.metrics import mean_squared_error\n",
    "from xgboost import XGBRegressor\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import LSTM, Dense\n",
    "\n",
    "# 设置中文显示\n",
    "plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "plt.rcParams['axes.unicode_minus'] = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据准备与特征工程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_data(stock_code='600519', start_date='20100101', end_date='20231231'):\n",
    "    \"\"\"\n",
    "    准备股票数据并生成特征\n",
    "    \n",
    "    Args:\n",
    "        stock_code: 股票代码\n",
    "        \n",
    "    Returns:\n",
    "        tuple: (特征DataFrame, 目标变量Series)\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # 获取股票数据\n",
    "        df = ak.stock_zh_a_hist(symbol=stock_code, period=\"daily\", \n",
    "                              start_date=start_date, end_date=end_date, adjust=\"qfq\")\n",
    "        df['日期'] = pd.to_datetime(df['日期'])\n",
    "        df = df.set_index('日期').sort_index()\n",
    "        \n",
    "        # 计算技术指标\n",
    "        df['5日均线'] = df['收盘'].rolling(5).mean()\n",
    "        df['20日均线'] = df['收盘'].rolling(20).mean()\n",
    "        df['60日均线'] = df['收盘'].rolling(60).mean()\n",
    "        \n",
    "        # 计算收益率\n",
    "        df['5日收益率'] = df['收盘'].pct_change(5)\n",
    "        df['20日收益率'] = df['收盘'].pct_change(20)\n",
    "        \n",
    "        # 计算波动率\n",
    "        df['5日波动率'] = np.log(df['收盘']/df['收盘'].shift(1)).rolling(5).std()\n",
    "        \n",
    "        # 目标变量: 未来5日收益率\n",
    "        target = df['收盘'].pct_change(5).shift(-5)\n",
    "        \n",
    "        # 删除缺失值\n",
    "        df = df.dropna()\n",
    "        target = target[df.index]\n",
    "        \n",
    "        return df, target\n",
    "    except Exception as e:\n",
    "        print(f\"数据准备失败: {e}\")\n",
    "        return pd.DataFrame(), pd.Series()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## XGBoost模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_xgboost(X_train, y_train, X_test, y_test):\n",
    "    \"\"\"\n",
    "    训练XGBoost模型\n",
    "    \n",
    "    Args:\n",
    "        X_train: 训练集特征\n",
    "        y_train: 训练集目标\n",
    "        X_test: 测试集特征\n",
    "        y_test: 测试集目标\n",
    "        \n",
    "    Returns:\n",
    "        tuple: (模型, 测试集预测结果)\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # 训练模型\n",
    "        model = XGBRegressor(n_estimators=100, learning_rate=0.1, random_state=42)\n",
    "        model.fit(X_train, y_train)\n",
    "        \n",
    "        # 预测\n",
    "        y_pred = model.predict(X_test)\n",
    "        \n",
    "        # 评估\n",
    "        rmse = np.sqrt(mean_squared_error(y_test, y_pred))\n",
    "        print(f\"XGBoost RMSE: {rmse:.4f}\")\n",
    "        \n",
    "        return model, y_pred\n",
    "    except Exception as e:\n",
    "        print(f\"XGBoost训练失败: {e}\")\n",
    "        return None, None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LSTM模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_lstm_data(data, target, time_steps=30):\n",
    "    \"\"\"\n",
    "    准备LSTM数据\n",
    "    \n",
    "    Args:\n",
    "        data: 特征数据\n",
    "        target: 目标变量\n",
    "        time_steps: 时间步长\n",
    "        \n",
    "    Returns:\n",
    "        tuple: (特征数组, 目标数组)\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # 归一化\n",
    "        scaler = MinMaxScaler()\n",
    "        scaled_data = scaler.fit_transform(data)\n",
    "        \n",
    "        # 创建时间序列数据\n",
    "        X, y = [], []\n",
    "        for i in range(time_steps, len(data)):\n",
    "            X.append(scaled_data[i-time_steps:i])\n",
    "            y.append(target.iloc[i])\n",
    "            \n",
    "        return np.array(X), np.array(y), scaler\n",
    "    except Exception as e:\n",
    "        print(f\"LSTM数据准备失败: {e}\")\n",
    "        return None, None, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_lstm(X_train, y_train, X_test, y_test):\n",
    "    \"\"\"\n",
    "    训练LSTM模型\n",
    "    \n",
    "    Args:\n",
    "        X_train: 训练集特征\n",
    "        y_train: 训练集目标\n",
    "        X_test: 测试集特征\n",
    "        y_test: 测试集目标\n",
    "        \n",
    "    Returns:\n",
    "        tuple: (模型, 测试集预测结果)\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # 构建模型\n",
    "        model = Sequential([\n",
    "            LSTM(50, return_sequences=True, input_shape=(X_train.shape[1], X_train.shape[2])),\n",
    "            LSTM(50),\n",
    "            Dense(1)\n",
    "        ])\n",
    "        \n",
    "        model.compile(optimizer='adam', loss='mse')\n",
    "        \n",
    "        # 训练模型\n",
    "        model.fit(X_train, y_train, epochs=20, batch_size=32, verbose=1)\n",
    "        \n",
    "        # 预测\n",
    "        y_pred = model.predict(X_test).flatten()\n",
    "        \n",
    "        # 评估\n",
    "        rmse = np.sqrt(mean_squared_error(y_test, y_pred))\n",
    "        print(f\"LSTM RMSE: {rmse:.4f}\")\n",
    "        \n",
    "        return model, y_pred\n",
    "    except Exception as e:\n",
    "        print(f\"LSTM训练失败: {e}\")\n",
    "        return None, None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 主分析流程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 准备数据\n",
    "features, target = prepare_data()\n",
    "\n",
    "# 划分训练测试集\n",
    "split = int(len(features)*0.8)\n",
    "X_train, X_test = features.iloc[:split], features.iloc[split:]\n",
    "y_train, y_test = target.iloc[:split], target.iloc[split:]\n",
    "\n",
    "# 训练XGBoost\n",
    "xgb_model, xgb_pred = train_xgboost(X_train, y_train, X_test, y_test)\n",
    "\n",
    "# 训练LSTM\n",
    "X_lstm, y_lstm, scaler = create_lstm_data(features, target)\n",
    "X_train_lstm, X_test_lstm = X_lstm[:split], X_lstm[split:]\n",
    "y_train_lstm, y_test_lstm = y_lstm[:split], y_lstm[split:]\n",
    "lstm_model, lstm_pred = train_lstm(X_train_lstm, y_train_lstm, X_test_lstm, y_test_lstm)\n",
    "\n",
    "# 绘制预测结果\n",
    "plt.figure(figsize=(12, 6))\n",
    "plt.plot(y_test.values, label='实际值')\n",
    "plt.plot(xgb_pred, label='XGBoost预测')\n",
    "plt.plot(lstm_pred, label='LSTM预测')\n",
    "plt.title('机器学习模型预测结果对比')\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}